#include "head.h"
#include <iostream>
#include <iomanip>
#include "mpi.h"

void read_matrix_01(){
    minm_cst = 300;
    fstream fin(FILENAME);
    int tmp;
    fin >> tmp;
    fin >> tmp;
    numInput = tmp;
    fin >> tmp;
    numTarget = tmp;
    numOutput = tmp;
    memset(output, 0, sizeof(output));
    for(int i = 0; i < numTarget; i++){
        for(int j = 0; j < numInput; j++){
            fin >> tmp;
            output[i] |= tmp << j;
        }
        index_output[output[i]] = i;
    }
    
    // for(int i = 0; i < numTarget; i++)
    //     cout << output[i] << ' ';
    // for(auto& [o, i]: id){
    //     cout << i << ' ';
    //     printBinary(o);
    // }
        
    fin.close();
}

void loop_init_state(vector<int> &s, int k) {
    for (int i = 0; i < k; i ++) {
        s[i] = i;
    }
}

void loop_init_state(vector<int> &s, int k, int st) {
    for (int i = 0; i < k; i ++) {
        s[i] = st + i;
    }
}

//  转到下一个状态
bool loop_next_state(vector<int> &s, int k, int n) {
    int i, j;

    i = k - 1;
    while ((i >= 0) && (s[i] == (n + i - k))) i --;   //  把大于等于0的判断放在前面比较安全
    if (i < 0) return false;
    //  i是从后往前第一个s.t. s[i] != n + i - k的元素
    //  第i个元素+1，然后i+i个元素开始，每个元素变成前一个元素+1，这样第二个相当于+2，第三个相当于+3
    s[i] ++;
    for (j = i + 1; j < k; j ++) {
        s[j] = s[j-1] + 1;
    }

    return true;
}

bool loop_next_state(vector<int> &s, int k, int n, int st) {
    int i, j;

    i = k - 1;
    while ((i >= 0) && (s[i] == (st + n + i - k))) i --;   //  把大于等于0的判断放在前面比较安全
    if (i < 0) return false;
    //  i是从后往前第一个s.t. s[i] != n + i - k的元素
    //  第i个元素+1，然后i+i个元素开始，每个元素变成前一个元素+1，这样第二个相当于+2，第三个相当于+3
    s[i] ++;
    for (j = i + 1; j < k; j ++) {
        s[j] = s[j-1] + 1;
    }

    return true;
}


void initial_base(vector<uint_t>& base, vector<int>& d, int numInput){
    for(int i = 0; i < numInput; i++){
        base.push_back(1 << i);
        d.push_back(0);
    }
}


bool choose_comb(uint_t T, vector<uint_t>& base, vector<int>& base_depth, int K, int S, int L)
{
	int base_sz = base.size();
    if ((base_sz -  S) < K) return false; //not enough base elements
    if (L < 1) return false;
    if (K==0) return false; //this is probably not reached
    if (K==1)
    {
      for (int i=S; i < base_sz; i++) if ((T == base[i]) && (1 << base_depth[i]) <= L) return true;
      return false;
    }

    //consider those sums containing Base[S]
    if (choose_comb(T^base[S],base, base_depth, K-1, S+1, L - (1 << base_depth[S]))) return true;
    //consider those sums not containing Base[S]
    if (choose_comb(T, base, base_depth,  K, S+1, L)) return true;
    //not found
    return false;
}







void get_choice_1_step(vector<uint_t>& choices, vector<uint_t>& base, vector<int>& base_depth, unordered_map<uint_t, int>& dist_table, vector<int>& dist, vector<vector<int>>& dists, vector<int>& new_depth, uint_t* target, int* bound_d){
    int sz = base.size();
    //base中选两个
    vector<uint_t> t_i_j; //base 中第 i ^ 第 j 的结果
    vector<pii> index_t_i_j;

    uint_t tmp;
    int new_d;
    for(int i = 0; i < sz; i++){
        if(base_depth[i] >= 3)
            continue;//顺序？
        for(int j = i + 1; j < sz; j++){
            tmp = base[i] ^ base[j];
            new_d = max(base_depth[i], base_depth[j]) + 1;
            if(new_d >= mx_depth)
                continue;
            if(base_depth[j] >= mx_depth - 1)
                continue;

            if(d(tmp) < new_d){
                bool min_source = true;
                for(int p = 0; p < sz; p++){
                    for(int q = p + 1; q < sz; q++){
                        if((base[p] ^ base[q]) == tmp && max(base_depth[p], base_depth[q]) + 1 < new_d){
                            min_source = false;
                            break;
                        }
                    }
                    if(!min_source)
                        break;
                }
                if(!min_source)
                    continue;
            }
            bool ok = true;
            for(int k = 0; k < sz; k++)
                if(tmp == base[k] && new_d >= base_depth[k]){
                    ok = false;
                    break;
                }
            if(!ok)
                continue;
            t_i_j.push_back(base[i] ^ base[j]);
            index_t_i_j.emplace_back(i, j);
        }
    }

    
    int min_d = 10000;//maxnum
    int min_eu_norm = 0;
    bool ok_mindist;
    vector<int> selected_index;


    int minmdist = INT_MAX;
    for(int i = 0; i < numTarget; i++){
        if(dist[i] == 1)
            continue;
        minmdist = min(minmdist, dist[i]);
    }

    int comb_sz = t_i_j.size();
    for(int i = 0; i < comb_sz; i++){
        int d = 0;
        int eu_norm = 0;
        uint_t mask = 0;
        int new_base_depth = max(base_depth[index_t_i_j[i].first], base_depth[index_t_i_j[i].second]) + 1;

        vector<int> tmpdist(dist);
        ok_mindist = false;
        for(int j = 0; j < numTarget; j++){
            if(dist[j] == 1){
                d += 1;
                eu_norm += 1;
                continue;
            }
            tmp = target[j] ^ t_i_j[i];
            if(choose_comb(tmp, base, base_depth, dist[j] - 2, 0, (1 << bound_d[j]) - (1 << new_base_depth))){
                tmpdist[j]--;
                if(dist[j] == minmdist)
                    ok_mindist = true;
            }
            d += tmpdist[j];
            eu_norm += tmpdist[j] * tmpdist[j];
        }

        if(!ok_mindist)
            continue;
        if(d < min_d || d == min_d && eu_norm < min_eu_norm){
            min_d = d;
            min_eu_norm = eu_norm;
            
            choices.clear();
            dists.clear();
            new_depth.clear();
            
            choices.push_back(t_i_j[i]);
            dists.emplace_back(tmpdist);
            new_depth.push_back(new_base_depth);
        }
        else if(d == min_d && eu_norm == min_eu_norm){
            choices.push_back(t_i_j[i]);
            dists.emplace_back(tmpdist);
            new_depth.push_back(new_base_depth);
        }
    }
}  

bool verify(vector<uint_t>& base, uint_t* target){
    bool exist[numTarget]{};
    for(int i = numInput; i < base.size(); i++){
        bool ok = false;
        for(int j = 0; j < i; j++)
            for(int k = j + 1; k < i; k++)
                if((base[j] ^ base[k]) == base[i]){
                    ok = true;
                }
        if(ok == false)
            return false;
        for(int j = 0; j < numTarget; j++)
            if(base[i] == target[j]){
                exist[j] = true;
                break;
            }
        
    }
    for(int i = 0; i < numTarget; i++)
        if(!exist[i])
            return false;
    return true;
}

bool verify_output(vector<uint_t>& base){
    bool exist[mxcst]{};
    int base_depth[mxcst]{};
    for(int i = numInput; i < base.size(); i++){
        bool ok = false;
	    int min_depth = mx_depth + 1;
        for(int j = 0; j < i; j++)
            for(int k = j + 1; k < i; k++)
                if((base[j] ^ base[k]) == base[i]){
                    ok = true;
		            min_depth = min(min_depth, max(base_depth[j], base_depth[k]) + 1);
                }
        if(ok == false || min_depth > mx_depth)
            return false;
        base_depth[i] = min_depth;
        for(int j = 0; j < numOutput; j++)
            if(base[i] == output[j]){
                exist[j] = true;
                break;
            }
        
    }
    for(int i = 0; i < numOutput; i++)
        if(!exist[i])
            return false;
    return true;
}


void bp(int useless, vector<uint_t>& base, vector<int>& base_depth, unordered_map<uint_t, int>& dist_table, vector<int>& dist, uint_t* target, int* bound_d, vector<tuuu>& split){
    int pres = 0;
    
    while(true){
        int s = 0;
        int best_j, best_k, reach_i, best_depth = INT_MAX;
        bool ok = false;
        for(int i = 0; i < numTarget; i++){
            if(dist[i] == 1){
                s += 1;
            }
            else if(dist[i] == 2){
                for(int j = 0; j < base.size(); j++){
                    for(int k = j + 1; k < base.size(); k++)
                        if((base[j] ^ base[k]) == target[i] && base_depth[j] < bound_d[i] && base_depth[k] < bound_d[i]){
                            reach_i = i;
                            if(max(base_depth[j], base_depth[k]) + 1 < best_depth){
                                best_depth = max(base_depth[j], base_depth[k]) + 1;
                                best_j = j;
                                best_k = k;
                            }
                            ok = true;
                        }
                }
                if(ok)
        		    break;

            }
            else{
                s += dist[i];
            }
          
        }
        if(ok){
            dist[reach_i] = 1;
            int new_depth = max(base_depth[best_j], base_depth[best_k]) + 1;
            for(int l = 0; l < numTarget; l++){
                if(dist[l] == 1)
                    continue;
                if(choose_comb(target[l] ^ target[reach_i], base, base_depth, dist[l] - 2, 0, (1 << bound_d[l]) - (1 << new_depth)))
                    dist[l]--;
            }
            base.emplace_back(target[reach_i]);
            base_depth.emplace_back(new_depth);                  
            continue;
        }
            
       
        if(s == pres){
            cout << "already " << base.size() - numInput << endl;
            return;
        }

        if(s == numTarget){


            tot_time++;
            if(tot_time % 200 == 0)
                cout << RANK << CHOICE << " minm now " << minm_cst << " OPT " << OPT << endl;
            if(verify(base, target)){
                
                int forward_basesz = base.size();
                for(int t = split.size() - 1; t >= 0; t--){
                    base.push_back(get<0>(split[t]));
                    base_depth.push_back(mx_depth);
                }
                if(verify_output(base)){
                    int xor_cnt = base.size() - numInput;

                    if(xor_cnt < OPT){
                        string record_better = "results/" + string(OUTPUT_FILENAME) + RANK + ".txt";
                        ofstream fout(record_better, ios::app);
                        
                        fout << CHOICE << "use" << xor_cnt << " xors, split size " << split.size() << endl;
                        for(int t = 0; t < base.size(); t++){
                            fout << base[t] << ' ';
                        }
                        
                        fout << endl << " cnt1 " << cnt1 << " cnt2 " << cnt2 << " cnt3 " << cnt3 << " cnt4 " << cnt4 << " cntrandom " << cntrandom  << endl;
                        
                        for(int i = numInput; i < base.size(); i++){
                            bool ok = false;
                            for(int j = 0; j < i; j++){
                                for(int k = j + 1; k < i; k++){
                                    if((max(base_depth[j], base_depth[k]) + 1 <= base_depth[i]) && (base[j] ^ base[k]) == base[i]){
                                        int flag = -1;
                                        for(int l = 0; l < numOutput; l++){
                                            if(base[i] == output[l]){
                                                flag = l;
                                                break;
                                            }
                                        }
                                        fout << "x" << i << " x" << j << " x" << k;
                                        if(flag >= 0)
                                            fout << " y" << flag;

                                        fout << " depth " << max(base_depth[j], base_depth[k]) + 1 << endl;
                                        ok = true;
                                        break;
                                    }
                                }
                                if(ok)
                                    break;
                            }


                        }
                        
                        fout << endl << "xors" << xor_cnt << endl << endl;
                        fout.close();
                    }

                    if(xor_cnt < minm_cst){
                        minm_cst = xor_cnt;
                        cout << RANK << CHOICE << "new minm xorcnt " << minm_cst << " , split size " << split.size() << " OPT " << OPT << endl;
                        cout << RANK << " cnt1 " << cnt1 << " cnt2 " << cnt2 << " cnt3 " << cnt3 << " cnt4 " << cnt4 << " cntrandom " << cntrandom  << endl;
                    }
                    else if(xor_cnt == minm_cst){
                        ;
                    }
                    
                }
                else
                        cout << "false" << endl;

            }    
            else
                cout << "false" << endl;

            return;
        }
        
            vector<uint_t> choices;
            vector<vector<int>> dists;
            vector<int> new_depth;
            get_choice_1_step(choices, base, base_depth, dist_table, dist, dists, new_depth, target, bound_d);

            if(choices.size() == 0){
                cout << "no choice " << endl;
        	    return;
            }

            int i = rand() % choices.size();

            dist = dists[i];
            base.push_back(choices[i]);
            base_depth.push_back(new_depth[i]);
        pres = s;    
    }
    

}


uint_t random_split_once(uint_t w, int s){
    int weight_w = __builtin_popcount(w);
    int maxm = 1 << (s - 1);
    int minm = weight_w - maxm;
    int rd_w = minm == maxm? minm: minm + rand() % (maxm - minm);
    int chosen = 0, tot_need = rd_w;
    uint_t mask = 0;
    uint_t ret = 0;
    for(int i = 0; i < weight_w; i++){
        if(rand() / (double)RAND_MAX < (double)(tot_need - chosen) / (weight_w - i)){
            //下一个选取概率为剩下要选个数 / 剩下可选个数
            ret |= w & -w;
            chosen++;
            if(chosen == tot_need)
                break;
        }
        w &= w - 1;
    }
    return ret;
}


vector<tuuu> get_split_seq(vector<tuuu> split, unordered_set<uint> W, unordered_set<uint> P, int s, int recur_depth){//rule3,5random
    unordered_set<uint> originalP = P;
    int split_bound = numOutput - originalP.size();
    int sz2, sz3, sz4;
    cnt1 = 0, cnt2 = 0, cnt3 = 0, cnt4 = 0, cntrandom = 0;
    uint w1, w2, w12, q1, q2, w;

    while(s > 1){
        while(W.size()){
            if(split.size() < min(9, split_bound)){
                vector<uint_t> targets;
                vector<int> bound_ds;
                for(uint x: W){
                    targets.push_back(x);
                    bound_ds.push_back(s);
                }
                for(uint x: P){
                    targets.push_back(x);
                    if(originalP.count(x))
                        bound_ds.push_back(s);
                    else
                        bound_ds.push_back(s - 1);
                }
                bp_with_target(targets, split, bound_ds);                
            }
            else return {};



            sz2 = 0, sz3 = 0, sz4 = 0;
            bool rmved = false;
            for(auto ite = W.begin(); ite != W.end();){
                rmved = false;
                uint w = *ite;
                for(uint p: P){ 
                    q1 = w ^ p;
                    if(P.count(q1)){
                        rmved = true;
                        ite = W.erase(ite);
                        originalP.erase(q1);
                        originalP.erase(p);
                        split.emplace_back(w, q1, p);
                        break;
                    }
                }
                if(rmved){
                    cnt1++;
                    continue;
                }
                ite = next(ite);
            }
            


            

            //rule4 2to3  || 
            

            for(auto it = W.begin(); it != W.end(); it++){
                w1 = *it;
                for(uint p: P){
                    q1 = w1 ^ p;
                    if(d(q1) < s){
                        rule3[sz3++] = make_pair(p, q1);
                    }
                }
            }
            
            int rand_3_5 = rand() % 2;
            if(rand_3_5 == 0){
                
                if(sz3){
                    cnt3++;
                    int rd3 = rand() % sz3;
                    auto [p, q] = rule3[rd3];
                    w1 = p ^ q;
                    if(!P.count(q) && d(q) > 0)
                        P.insert(q);
                    originalP.erase(p);
                    originalP.erase(q);

                    split.emplace_back(w1, p, q);
                    W.erase(w1);
                    continue;
                }
                
            
                //random split once
                cntrandom++;
                

                auto ite = W.begin();
                int rd = rand() % W.size();
                while(rd--)
                    ite = next(ite);
                uint w = *ite;
                W.erase(ite);
                
                q1 = random_split_once(w, s);
                q2 = w ^ q1;
                if(!P.count(q1) && d(q1))
                    P.insert(q1);
                if(!P.count(q2) && d(q2))
                    P.insert(q2);
                originalP.erase(q1);
                originalP.erase(q2);
                split.emplace_back(w, q1, q2);    
            
                
            }
            else{
                //random split once
                cntrandom++;
                

                auto ite = W.begin();
                int rd = rand() % W.size();
                while(rd--)
                    ite = next(ite);
                uint w = *ite;
                W.erase(ite);
                
                q1 = random_split_once(w, s);
                q2 = w ^ q1;
                if(!P.count(q1) && d(q1))
                    P.insert(q1);
                if(!P.count(q2) && d(q2))
                    P.insert(q2);
                originalP.erase(q1);
                originalP.erase(q2);
                split.emplace_back(w, q1, q2);                    
            }

            
            
        }
        s--;
        for(auto ite = P.begin(); ite != P.end();){
            uint p = *ite;
            if(d(p) == s){
                W.insert(p);
                ite = P.erase(ite);
            }
            else
                ite++;
        }
    }

    for(uint w: W){
        q1 = w & (w - 1);
        q2 = w ^ q1;
        split.emplace_back(w, q1, q2);
    }
    // reverse(split.begin(), split.end());
    return split;
}


void run_mpi(int argc, char** argv){

    MPI_Init(&argc, &argv);
    int world_size;
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);
    int world_rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
    string mx_duration(argv[1]);
    int mx_dr = stoi(mx_duration);
    global_mx_dr = mx_dr;
    int a;
    clock_t st, en;
    st = clock();
    global_st = clock();
    srand(time(NULL) + (unsigned long)(&a));

    vector<tuuu> split;
    unordered_set<uint> W;
    unordered_set<uint> P;
    read_matrix_01();
    RANK = "_" + to_string(world_rank) + "_";
    cout << "choice " << CHOICE << " world rank " << world_rank << " duration " << mx_dr << endl; 
    int s = 0;
    for(int i = 0; i < numTarget; i++){
        s = max(s, d(output[i]));
        mx_depth = max(mx_depth, s);
    }
    for(int i = 0; i < numTarget; i++){
        if(d(output[i]) == s)
            W.insert(output[i]);
        else
            P.insert(output[i]);    
    }
    while(true){
        en = clock();
        int duration = (double)(en - st) / CLOCKS_PER_SEC;
	// cout << "thread " << to_string(world_rank + st_ID) << " duration " << duration << endl; 
        if(duration > mx_dr)
            break;
        auto res = get_split_seq(split, W, P, s, MX_RECUR_DEPTH);
    }

    en = clock();
    int duration = (double)(en - st) / CLOCKS_PER_SEC;
    cout << RANK << " use " << duration << "seconds, minm" << minm_cst << "ite_cnt" << ite_cnt << endl;
    MPI_Finalize();
}


void bp_with_target(vector<uint_t>& targets, vector<tuuu>& split, vector<int>& bound_ds){
    //处理得到target
    uint_t target[mxcst];
    int bound_d[mxcst];
    // int a;
    // clock_t st, en;
    // st = clock();
    // srand(time(NULL) + (unsigned long)(&a));
    numTarget = targets.size();
    for(int i = 0; i < numTarget; i++)
        target[i] = targets[i], bound_d[i] = bound_ds[i];



    // srand((unsigned long)(&a) + time(NULL));


    minm_cst = min(minm_cst, 300);
    
    int repeat = 10;
    while(repeat--){
	global_en = clock();
	int global_duration = (double)(global_en - global_st) / CLOCKS_PER_SEC;
	if(global_duration > global_mx_dr)
	    break;
	ite_cnt++;

        vector<uint_t> base;
        vector<int> base_depth;
        base.reserve(mxcst);
        base_depth.reserve(mxcst);
        initial_base(base, base_depth, numInput);


        vector<int> dist(numTarget);//项数，异或数+1

        //initial dist
        for(int i = 0; i < numTarget; i++){
            int w = __builtin_popcount(target[i]);
            dist[i] = w;
        }


        unordered_map<uint_t, int> dist_table;
        
        bp(minm_cst, base, base_depth, dist_table, dist, target, bound_d, split);

    }

}

void printbinary(uint_t x){
    for(int i = 0; i < numInput; i++)
        if((x >> i) & 1)
            cout << "1 ";
        else
            cout << "0 ";
    cout << endl;
}
